/*
 * Copyright (c) 2022. China Mobile (SuZhou) Software Technology Co.,Ltd. All rights reserved.
 * Lakehouse is licensed under Mulan PSL v2.
 * You can use this software according to the terms and conditions of the Mulan PSL v2.
 * You may obtain a copy of Mulan PSL v2 at:
 *          http://license.coscl.org.cn/MulanPSL2
 * THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND,
 * EITHER EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT,
 * MERCHANTABILITY OR FIT FOR A PARTICULAR PURPOSE.
 * See the Mulan PSL v2 for more details.
 */

package com.chinamobile.cmss.lakehouse.service.hive;

import java.sql.Connection;
import java.sql.SQLException;

import org.apache.commons.pool2.KeyedPooledObjectFactory;
import org.apache.commons.pool2.PooledObject;
import org.apache.commons.pool2.impl.DefaultPooledObject;

public final class HiveConnectionFactory implements KeyedPooledObjectFactory<String, Connection> {

    private final HiveConnector hiveConnector;

    private static final boolean DEFAULT_AUTO_COMMIT = true;

    public HiveConnectionFactory(HiveConnector hiveConnector) {
        this.hiveConnector = hiveConnector;
    }

    @Override
    public PooledObject<Connection> makeObject(String s) throws Exception {
        Connection connection = hiveConnector.connect(s);
        return new DefaultPooledObject<>(connection);
    }

    @Override
    public void destroyObject(String s, PooledObject<Connection> pooledObject) throws Exception {
        pooledObject.getObject().close();
    }

    @Override
    public boolean validateObject(String s, PooledObject<Connection> pooledObject) {
        try {
            final int timeout = 60 * 1000;
            return pooledObject.getObject().isValid(timeout);
        } catch (SQLException sqlException) {
            throw new RuntimeException(sqlException);
        }
    }

    @Override
    public void activateObject(String s, PooledObject<Connection> pooledObject) throws Exception {
        Connection conn = pooledObject.getObject();
        // 对象被借出，需要进行初始化，将其 autoCommit进行设置
        if (conn.getAutoCommit() != DEFAULT_AUTO_COMMIT) {
            conn.setAutoCommit(DEFAULT_AUTO_COMMIT);
        }
    }

    @Override
    public void passivateObject(String s, PooledObject<Connection> pooledObject) throws Exception {
        Connection conn = pooledObject.getObject();
        if (!conn.getAutoCommit() && !conn.isReadOnly()) {
            conn.rollback();
        }
        conn.clearWarnings();
        if (!conn.getAutoCommit()) {
            conn.setAutoCommit(true);
        }
    }
}